# 提取和输出结构化数据

import os
from typing import List

from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnablePassthrough
from langchain_openai.chat_models.base import BaseChatOpenAI
from pydantic.v1 import BaseModel, Field

# 本地有clash代理，所以配置一下，不然一些依赖下载不下来（chroma）,或者将代理关闭后下载
os.environ['http_proxy'] = '127.0.0.1:7890'
os.environ['https_proxy'] = '127.0.0.1:7890'

os.environ["LANGCHAIN_TRACING_V3"] = "true"
os.environ["LANGCHAIN_API_KEY"] = "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"
os.environ["LANGCHAIN_PROJECT"] = "langchain-community-demo"

os.environ["OPENAI_API_KEY"] = "sk-1dd16a258a73428d910d38c782e1c94f"

# deepseek-reasoner : DeepSeek-R1
# deepseek-chat : DeepSeek-V3
model_name = "deepseek-chat"
deepseek_api_key = "sk-1dd16a258a73428d910d38c782e1c94f"

model = BaseChatOpenAI(
    model=model_name,
    openai_api_key=deepseek_api_key,
    openai_api_base='https://api.deepseek.com',
    max_tokens=1024,
    streaming=True
)


# pydantic: 处理数据、验证数据、定义数据格式、虚拟化和反虚拟化、类型转换等
class Person(BaseModel):
    name: str = Field(default=None, description="姓名")
    age: int = Field(default=None, description="年龄")
    hight: float = Field(default=None, description="身高")


class ManyPeople(BaseModel):
    people: List[Person] = Field(default=None, description="人")


# 定义一个提示词模版
prompt = ChatPromptTemplate.from_messages([
    (
        "system",
        "你是一个专业提取算法的AI，你需要根据用户问题，从未结构化文本中提取结构化数据，如果你不知道属性的值，返回该属性值为null。"
    ),
    # 历史记录
    # MessagesPlaceholder("examples"),
    ("human", "{text}")
])

# 只能提取一个结构化数据
# chain = {'text': RunnablePassthrough()} | prompt | model.with_structured_output(schema=Person)

# 可以提取多个结构化数据
chain = {'text': RunnablePassthrough()} | prompt | model.with_structured_output(schema=ManyPeople)

text = "朝我迎面走来一个女生，大概有18岁的样子，身体估计有一米七，旁边还有一个男生，是我的朋友，叫王三，20岁了，身高有一米八。"
resp = chain.invoke({"text": text})
print(resp)
